// Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use arrow::datatypes::ToByteSlice;
use bytes::{BufMut, Bytes, BytesMut};
use dora_node_api::{self, DoraNode, Event, EventStream};
use eyre::eyre;
use rerun::{components::TensorData, external::{anyhow::Context, arrow2, clap::{self, Parser}, log::{error, info}, re_log::{self}}, TensorBuffer};
use rerun_tools::{dora_data::DataKind, TensorBufferType};
use std::path::PathBuf;

#[derive(Debug, Parser)]
#[clap(author, version, about)]
struct Args {
    /// The path to image file
    image_path: PathBuf,
}
fn tensor_data_to_bytes(tensor_data: TensorData) -> Bytes {
    let nbytes = size_of::<u64>()
        + size_of::<u64>()
        + tensor_data.shape.len() * size_of::<u64>()
        + size_of::<u8>()
        + tensor_data.buffer.size_in_bytes();
    let mut bytes = BytesMut::with_capacity(nbytes);
    bytes.put_u64(nbytes as u64);
    bytes.put_u64(tensor_data.shape.len().try_into().unwrap());
    for dimension in &tensor_data.shape {
        bytes.put_u64(dimension.size);
    }
    let tensor_buffer_type = TensorBufferType::try_from(&tensor_data.buffer)
        .expect("read tensor buffer type failed");
    let tensor_buffer_type = tensor_buffer_type as u8;
    bytes.put_u8(tensor_buffer_type);
    bytes.put_slice(match &tensor_data.buffer {
        TensorBuffer::U8(arrow_buffer) => arrow_buffer.to_byte_slice(),
        TensorBuffer::U16(arrow_buffer) => arrow_buffer.to_byte_slice(),
        TensorBuffer::U32(arrow_buffer) => arrow_buffer.to_byte_slice(),
        TensorBuffer::U64(arrow_buffer) => arrow_buffer.to_byte_slice(),
        TensorBuffer::I8(arrow_buffer) => arrow_buffer.to_byte_slice(),
        TensorBuffer::I16(arrow_buffer) => arrow_buffer.to_byte_slice(),
        TensorBuffer::I32(arrow_buffer) => arrow_buffer.to_byte_slice(),
        TensorBuffer::I64(arrow_buffer) => arrow_buffer.to_byte_slice(),
        TensorBuffer::F16(arrow_buffer) => unsafe {
            std::slice::from_raw_parts(
                arrow_buffer.as_ptr() as *const u8,
                arrow_buffer.len() * size_of::<arrow2::types::f16>(),
            )
        },
        TensorBuffer::F32(arrow_buffer) => arrow_buffer.to_byte_slice(),
        TensorBuffer::F64(arrow_buffer) => arrow_buffer.to_byte_slice(),
        TensorBuffer::Jpeg(arrow_buffer) => arrow_buffer.to_byte_slice(),
        TensorBuffer::Nv12(arrow_buffer) => arrow_buffer.to_byte_slice(),
        TensorBuffer::Yuy2(arrow_buffer) => arrow_buffer.to_byte_slice(),
    });
    bytes.freeze()
}

fn run(mut node: DoraNode, mut events: EventStream, path: PathBuf) -> eyre::Result<()> {
    let mut has_sent_images = true;
    let image = rerun::Image::from_file_path(path)
        .context("failed to load image");
    let tensor_data = image.unwrap().data;
    loop {
        let event = match events.recv() {
            Some(input) => input,
            None => break,
        };
        match event {
            Event::Stop(_) => {
                info!("dora stop");
                return Ok(());
            },
            Event::Reload { operator_id } => {
                return Err(eyre!("unhandled dora reload event: {:?}", operator_id));
            },
            Event::Input {
                id,
                metadata: _,
                data: _,
            } => match id.as_str() {
                "tick" => {
                    if has_sent_images {
                        let array = tensor_data_to_bytes(tensor_data.clone());
                        node.send_output_bytes(
                            String::from(DataKind::Image.as_output_id()).into(),
                            Default::default(),
                            array.len(),
                            &array,
                        )?;
                        has_sent_images = false;
                    }
                }
                other => {
                    error!("unexpected input: {}", other);
                }
            },
            Event::InputClosed { id } => {
                info!("dora input {} closed", id);
                return Ok(());
            },
            Event::Error(msg) => {
                return Err(eyre!("received dora error event: {}", msg));
            },
            _ => {
                error!("unknown dora event: {:?}", event);
            }
        }
    }
    Ok(())
}

#[tokio::main]
async fn main() -> eyre::Result<()> {
    re_log::setup_logging();
    let args = Args::try_parse()?;
    let (node, events) = DoraNode::init_from_env()?;
    run(node, events, args.image_path)?;
    Ok(())
}
